import json

import numpy as np

from analyse.data_model import Job
from model import EvaluateNet
import torch
import matplotlib.pyplot as plt


def load_data(data_path_):
    with open(data_path_, encoding='utf8', mode='r') as f:
        json_data = json.load(f)
        train_data_ = list()
        train_label_ = list()
        for item in json_data:
            job = Job.from_dict(item)
            train_data_.append(job.get_job_features())
            train_label_.append(job.generate_label())
        return np.array(train_data_), np.array(train_label_)


def standardization(data):
    """
    标准化
    :param data:
    :return:
    """
    mu = np.mean(data, axis=0)
    sigma = np.std(data, axis=0)
    return (data - mu) / sigma, mu, sigma


def normalization(data):
    """
    归一化
    :param data:
    :return:
    """
    _range = np.max(data) - np.min(data)
    return (data - np.min(data)) / _range


def draw_train_curve(loss_list, label, pred):
    """
    训练曲线绘制
    :param loss_list:
    :param label:
    :param pred:
    :return:
    """
    plt.figure(1)
    plt.subplot(2, 1, 1)
    plt.title("Loss")
    plt.plot(list(range(len(loss_list))), [i.item() for i in loss_list])
    plt.subplot(2, 1, 2)
    # predict
    plt.title("cmp")
    label = [i.item() for i in label.reshape((-1))]
    pred = [i.item() for i in pred.reshape((-1))]
    plt.plot(list(range(len(label))), label, label='label')
    plt.plot(list(range(len(pred))), pred, label='predict')
    plt.legend()
    plt.show()


def train(
        train_data_,
        train_label_,
        learning_rate,
        train_epoch
):
    """
    训练
    :param train_data_:
    :param train_label_:
    :param learning_rate:
    :param train_epoch: 训练轮数
    :return:
    """
    x, mean_, std_ = standardization(train_data_)
    x = torch.tensor(x, dtype=torch.float)
    y = normalization(train_label_).reshape((-1, 1))
    y = torch.tensor(y, dtype=torch.float)

    model_ = EvaluateNet(learning_rate)

    loss_list = list()
    for t in range(train_epoch):
        y_pred = model_(x)
        loss = model_.criterion(y_pred, y)
        loss_list.append(loss.detach())
        model_.optimizer.zero_grad()
        loss.backward()
        model_.optimizer.step()
    y_predict = model_(x)
    draw_train_curve(loss_list, y, y_predict)
    return model_, mean_.tolist(), std_.tolist()


def save_model(model_: torch.nn.Module, mean_, std_, pth_path, std_path):
    torch.save(model_.state_dict(), pth_path)
    with open(std_path, encoding='utf8', mode='w') as f:
        json.dump({
            'stand': {
                'mean': mean_,
                'std': std_
            }
        }, f)


if __name__ == '__main__':
    data_path = "./data/job_train_data.json"
    train_data, train_label = load_data(data_path)
    model, mean, std = train(
        train_data,
        train_label,
        learning_rate=3e-4,
        train_epoch=100
    )
    save_model(
        model, mean, std,
        pth_path="./check_point/eval_net.pth",
        std_path="./check_point/eval_net.json"
    )
